import os
import requests
import json
from volcenginesdkarkruntime import Ark

class ImageGenerator:
    def __init__(self, api_key, base_url="https://ark.cn-beijing.volces.com/api/v3"):
        self.client = Ark(
            base_url=base_url,
            api_key=api_key,
        )
        self.base_url = base_url
    
    def generate_image(self, prompt):
        try:
            imagesResponse = self.client.images.generate(
                model="doubao-seedream-4-0-250828",
                prompt=prompt,
                size="2K",
                response_format="url",
                watermark=True
            )
            
            # 返回结果
            return {
                "url": imagesResponse.data[0].url,
                "size": imagesResponse.data[0].size
            }
        except Exception as e:
            raise Exception(f"图像生成失败: {str(e)}")
    
    def upload_image_to_get_url(self, image_path):
        try:
            # 读取图片文件
            with open(image_path, 'rb') as image_file:
                files = {'image': (os.path.basename(image_path), image_file, 'image/jpeg')}
                response = requests.post(
                    f"{self.base_url}/upload",
                    headers={"Authorization": f"Bearer {self.client.api_key}"},
                    files=files
                )
                response.raise_for_status()
                result = response.json()
                return result['data']['url']
        except Exception as e:
            raise Exception(f"图片上传失败: {str(e)}")
    
    def generate_image_from_image(self, prompt, image_url):
        try:
            imagesResponse = self.client.images.generate(
                model="doubao-seedream-4-0-250828",
                prompt=prompt,
                image=image_url,
                size="2K",
                response_format="url",
                watermark=True
            )
            
            # 返回结果
            return {
                "url": imagesResponse.data[0].url,
                "size": imagesResponse.data[0].size
            }
        except Exception as e:
            raise Exception(f"图像生成失败: {str(e)}")